from flask import Blueprint, jsonify, request
import pandas as pd
import numpy as np
from enum import Enum
from datetime import datetime

# 创建蓝图对象
back_trade = Blueprint('back', __name__)

# 工具函数
def get_data(codes, back_cycle='day', is_real=False, start_date=None):
    """获取交易数据"""
    relative_path = 'data/realgo_5m/' if is_real else 'data/5m/'
    result = {}
    
    for code in codes:
        try:
            df = pd.read_csv(f'{relative_path}{code}.csv')
            df['date'] = pd.to_datetime(df['date'])
            
            # 如果提供了开始日期，过滤数据
            if start_date:
                start_date_dt = pd.to_datetime(start_date)
                df = df[df['date'] >= start_date_dt]
                
            df.set_index('date', inplace=True)
            
            if back_cycle == 'day':
                filtered_df = df.between_time('14:50:00', '14:50:00')
            else:
                filtered_df = df.resample(back_cycle).agg({'close': 'first', 'code': 'first'})
            
            # 确保时间索引有效
            filtered_df = filtered_df[filtered_df.index.isin(df.index)]
            filtered_df.reset_index(inplace=True)
            filtered_df['date'] = filtered_df['date'].dt.strftime('%Y-%m-%d %H:%M:%S')
            
            result[code] = filtered_df
        except FileNotFoundError:
            print(f"警告：代码 {code} 的数据文件未找到")
            result[code] = pd.DataFrame()
    
    return result

def get_hs300_data(hs300_limit, back_cycle='day', start_date=None):
    """获取沪深300数据"""
    df = get_data(['000300.XSHG'], back_cycle, False, start_date)['000300.XSHG']
    map_cy = {
        'day': 1,
        '5T': 48,
        '15T': 16,
        '30T': 8,
        '60T': 4
    }
    df['MAshort'] = df['close'].rolling(window=map_cy[back_cycle] * hs300_limit).mean()
    df['condition'] = df['close'] > df['MAshort']
    return df

class TradingStrategy:
    """交易策略主类，封装所有策略相关逻辑"""
    
    class TradeType(Enum):
        DIFFERENT_MA = 1          # 双均线策略
        CLOSE_ABOVE_MA = 2        # 收盘价高于均线
        CLOSE_ABOVE_HIGH = 3      # 收盘价高于历史高点
        HIGH_LOW_STRATEGY = 4     # 高低点策略

    class ReturnType(Enum):
        PCT_CHANGE = 1            # 涨跌幅
        DEVIATION_FROM_MA = 2     # 偏离均线幅度

    def __init__(self, params):
        # 初始化策略参数
        self.params = params
        self.n = params['increaseday_range']
        self.long_window = params['highday_range']
        self.short_window = params['lowday_range']
        
        # 参数验证：确保窗口大小不为0
        if self.short_window <= 0:
            self.short_window = 1
            print(f"警告：lowday_range不能为0或负数，已设置为1")
        if self.long_window <= 0:
            self.long_window = 1
            print(f"警告：highday_range不能为0或负数，已设置为1")
        
        self.check_hs300 = params['ischeckhs300']
        self.sort_type = self.ReturnType(params['sortType'])
        self.trade_type = self.TradeType(params['tradeType'])
        
        # 初始化策略状态
        self.trade_log = []
        self.position = {}
        self.cumulative_return = 1.0
        self.cumulative_returns = []

    def _create_position(self, date, code, price):
        """创建持仓记录（私有方法）"""
        return {
            'date': date,
            'price': price,
            'code': code
        }

    def _log_trade(self, date, trade_type, code, price, desc=None):
        """记录交易日志（私有方法）"""
        trade_entry = {
            'date': date,
            'type': trade_type,
            'price': price,
            'code': code
        }
        if desc:
            trade_entry['desc'] = desc
        self.trade_log.append(trade_entry)

    def _calculate_returns(self, df):
        """计算收益率（私有方法）"""
        if self.sort_type == self.ReturnType.PCT_CHANGE:
            df['gap-turn'] = df['close'].pct_change(periods=self.n)
        elif self.sort_type == self.ReturnType.DEVIATION_FROM_MA:
            df['gap-turn'] = (df['close'] - df['close'].rolling(self.n).mean()) / df['close']
        
        df['cur-turn'] = df['close'].pct_change(periods=1)  # 每日涨跌幅
        return df

    def _calculate_moving_conditions(self, df):
        """计算均线条件（私有方法）"""
        if self.trade_type == self.TradeType.DIFFERENT_MA:
            df['MAshort'] = df['close'].rolling(window=self.short_window).mean()
            df['MAlong'] = df['close'].rolling(window=self.long_window).mean()
            df['condition'] = df['MAshort'] > df['MAlong']
        elif self.trade_type == self.TradeType.CLOSE_ABOVE_MA:
            df['MAlong'] = df['close'].rolling(window=self.long_window).mean()
            df['condition'] = df['close'] > df['MAlong']
        elif self.trade_type == self.TradeType.CLOSE_ABOVE_HIGH:
            df['MAlong'] = df['close'].shift(1).rolling(window=self.long_window).max()
            df['condition'] = df['close'] > df['MAlong']
        elif self.trade_type == self.TradeType.HIGH_LOW_STRATEGY:
            df['MAlong'] = df['close'].shift(1).rolling(window=self.long_window).max()
            df['MAshort'] = df['close'].shift(1).rolling(window=self.short_window).min()
            df['condition'] = df['close'] > df['MAlong']
            df['sell_condition'] = df['close'] < df['MAshort']
        
        return df

    def _execute_trade_decision(self, row, best_code, current_code):
        """执行交易决策（私有方法）"""
        is_empty = not bool(current_code)
        buy_condition = row.get(f'condition-{best_code}', False)
        sell_condition = row.get(f'sell_condition-{best_code}', False)
        hs300_condition = row.get('condition_hs300', True) if self.check_hs300 else True

        # 空仓时买入
        if is_empty and buy_condition and hs300_condition:
            self.position = self._create_position(
                date=row['date'],
                code=best_code,
                price=row[f'close-{best_code}']
            )
            self._log_trade(row['date'], 'buy', best_code, row[f'close-{best_code}'])

        # 持仓时卖出
        elif not is_empty:
            # 沪深300条件不满足
            if not hs300_condition:
                self._log_trade(row['date'], 'sell', current_code, row[f'close-{current_code}'], '沪深300低于30日线')
                self.position.clear()
                return

            # 最佳代码变更
            if best_code != current_code:
                self._log_trade(row['date'], 'sell', current_code, row[f'close-{current_code}'])
                self.position.clear()
                
                if buy_condition:  # 换仓买入
                    self.position = self._create_position(
                        date=row['date'],
                        code=best_code,
                        price=row[f'close-{best_code}']
                    )
                    self._log_trade(row['date'], 'buy', best_code, row[f'close-{best_code}'])

            # 策略卖出条件触发
            elif (self.trade_type == self.TradeType.HIGH_LOW_STRATEGY and sell_condition) or \
                 (self.trade_type != self.TradeType.HIGH_LOW_STRATEGY and not buy_condition):
                self._log_trade(row['date'], 'sell', current_code, row[f'close-{current_code}'])
                self.position.clear()

    def execute(self):
        """执行完整交易策略（主方法）"""
        # 获取标的数据
        df_map = get_data(
            self.params['codes'], 
            self.params['backCycle'], 
            self.params.get('isReal', False),
            self.params.get('startDate')
        )

        # 预处理数据
        processed_dfs = []
        for code in self.params['codes']:
            df = self._calculate_returns(df_map[code])
            df = self._calculate_moving_conditions(df)
            df = df.rename(columns={col: f'{col}-{code}' for col in df.columns if col != 'date'})
            processed_dfs.append(df)

        # 合并数据
        merged_df = pd.concat(processed_dfs, ignore_index=True) if len(processed_dfs) == 1 \
                    else pd.merge(processed_dfs[0], processed_dfs[1], on='date', how='outer')
        for df in processed_dfs[2:]:
            merged_df = pd.merge(merged_df, df, on='date', how='outer')

        # 处理最佳代码选择
        gap_turn_cols = [f'gap-turn-{code}' for code in self.params['codes']]
        merged_df['best_code'] = merged_df[gap_turn_cols].idxmax(axis=1, skipna=True)
        merged_df['best_code'] = merged_df['best_code'].str.extract(r'gap-turn-(.+)')  # 提取代码部分

        # 处理沪深300数据（如果需要）
        if self.check_hs300:
            hs300_df = self.params['hs300_data'].rename(columns={col: f'{col}_hs300' for col in self.params['hs300_data'].columns if col != 'date'})
            merged_df = pd.merge(merged_df, hs300_df, on='date', how='outer')

        # 执行交易循环
        for _, row in merged_df.iterrows():
            current_code = self.position.get('code')
            best_code = row['best_code']

            # 计算累计收益
            if current_code:
                daily_return = row.get(f'cur-turn-{current_code}', 0)
                self.cumulative_return *= (1 + daily_return)
            self.cumulative_returns.append(self.cumulative_return)

            # 跳过收盘时间和日内重复交易
            if row['date'][11:19] == '15:00:00':
                continue
            if current_code and row['date'][:10] == str(self.position['date'])[:10]:
                continue

            # 执行交易决策
            self._execute_trade_decision(row, best_code, current_code)
        print(f"gg--- Trade log: {self.trade_log} ---") # 调试输出
        print(f"gg--- Cumulative returns: {self.cumulative_returns} ---") # 调试输出
        return self.trade_log, pd.DataFrame({
            'date': merged_df['date'],
            'cumulative_return': self.cumulative_returns
        })

def calculate_drawdowns(returns_df):
    """计算回撤及持续时间"""
    # 确保数据按日期排序
    returns_df = returns_df.sort_values('date')
    
    # 将日期列转换为datetime类型
    returns_df['date'] = pd.to_datetime(returns_df['date'])
    
    # 计算累计收益的峰值
    returns_df['peak'] = returns_df['cumulative_return'].cummax()
    
    # 计算回撤
    returns_df['drawdown'] = (returns_df['cumulative_return'] / returns_df['peak']) - 1
    
    # 标记回撤开始和结束
    returns_df['is_drawdown'] = returns_df['drawdown'] < 0
    
    # 标记回撤开始和结束的索引
    drawdown_starts = []
    drawdown_ends = []
    drawdown_depths = []
    drawdown_durations = []
    
    in_drawdown = False
    start_idx = None
    peak_value = None
    
    for i, row in returns_df.iterrows():
        if not in_drawdown and row['drawdown'] < 0:
            # 回撤开始
            in_drawdown = True
            start_idx = i
            peak_value = row['peak']
        elif in_drawdown and (row['drawdown'] == 0 or i == len(returns_df) - 1):
            # 回撤结束或到达数据末尾
            in_drawdown = False
            drawdown_starts.append(returns_df.loc[start_idx, 'date'])
            drawdown_ends.append(row['date'])
            
            # 计算该回撤区间内的最大回撤深度
            drawdown_slice = returns_df.loc[start_idx:i]
            max_drawdown = drawdown_slice['drawdown'].min()
            drawdown_depths.append(max_drawdown)
            
            # 计算回撤持续天数
            duration = (row['date'] - returns_df.loc[start_idx, 'date']).days
            drawdown_durations.append(duration)
    
    # 创建回撤摘要DataFrame
    if drawdown_starts:
        drawdowns_df = pd.DataFrame({
            'start_date': drawdown_starts,
            'end_date': drawdown_ends,
            'depth': drawdown_depths,
            'duration': drawdown_durations
        })
        
        # 按回撤深度排序（从大到小，因为回撤是负值）
        drawdowns_df = drawdowns_df.sort_values('depth').reset_index(drop=True)
        
        # 取前5大回撤
        top_drawdowns = drawdowns_df.head(5).to_dict('records')
    else:
        top_drawdowns = []
    
    return top_drawdowns

def calculate_monthly_returns(returns_df):
    """计算每年每月的收益率"""
    # 确保数据按日期排序
    returns_df = returns_df.sort_values('date')
    
    # 将日期列转换为datetime类型
    returns_df['date'] = pd.to_datetime(returns_df['date'])
    
    # 提取年和月
    returns_df['year'] = returns_df['date'].dt.year
    returns_df['month'] = returns_df['date'].dt.month
    
    # 按年月分组，计算月度收益率
    monthly_returns = []
    
    # 获取所有年份和月份
    years = sorted(returns_df['year'].unique())
    
    for year in years:
        year_data = returns_df[returns_df['year'] == year]
        months = sorted(year_data['month'].unique())
        
        for month in months:
            month_data = year_data[year_data['month'] == month]
            
            if len(month_data) > 0:
                # 获取月初和月末的累计收益
                month_start = month_data.iloc[0]['cumulative_return']
                month_end = month_data.iloc[-1]['cumulative_return']
                
                # 计算月度收益率
                monthly_return = (month_end / month_start) - 1
                
                monthly_returns.append({
                    'year': int(year),
                    'month': int(month),
                    'return': float(monthly_return)
                })
    
    return monthly_returns

def calculate_yearly_returns(returns_df):
    """计算年度收益率"""
    # 确保数据按日期排序
    returns_df = returns_df.sort_values('date')
    
    # 将日期列转换为datetime类型
    returns_df['date'] = pd.to_datetime(returns_df['date'])
    
    # 提取年份和月份
    returns_df['year'] = returns_df['date'].dt.year
    returns_df['month'] = returns_df['date'].dt.month
    
    # 按年份分组，计算年度收益率
    yearly_returns = []
    years = sorted(returns_df['year'].unique())
    
    for year in years:
        year_data = returns_df[returns_df['year'] == year]
        
        if len(year_data) > 0:
            # 获取年初和年末的累计收益
            year_start = year_data.iloc[0]['cumulative_return']
            year_end = year_data.iloc[-1]['cumulative_return']
            
            # 计算年度收益率
            yearly_return = (year_end / year_start) - 1
            
            # 计算该年每个月的收益率，用于验证
            months_in_year = sorted(year_data['month'].unique())
            monthly_returns_in_year = []
            
            for month in months_in_year:
                month_data = year_data[year_data['month'] == month]
                if len(month_data) > 0:
                    month_start = month_data.iloc[0]['cumulative_return']
                    month_end = month_data.iloc[-1]['cumulative_return']
                    monthly_return = (month_end / month_start) - 1
                    monthly_returns_in_year.append(monthly_return)
            
            # 计算月度收益率累乘值 (1+r1)*(1+r2)*...*(1+rn)-1
            compound_monthly_return = 1.0
            for monthly_return in monthly_returns_in_year:
                compound_monthly_return *= (1 + monthly_return)
            compound_monthly_return -= 1
            
            yearly_returns.append({
                'year': int(year),
                'return': float(yearly_return),
                'compound_monthly_return': float(compound_monthly_return),
                'months_count': len(monthly_returns_in_year)
            })
    
    return yearly_returns

# 接口入口
@back_trade.route('/back', methods=['GET', 'POST'])
def backtesting_entry():
    """回测接口
    支持GET方法（使用默认参数）和POST方法（通过请求体接收参数）
    """
    # 默认参数
    default_params = {
        "ischeckhs300": False,
        "hs300_limit": 30,
        "backCycle": "day",
        "codes": ["000016.XSHG", "000852.XSHG", "399673.XSHE", "513100.XSHG", "518880.XSHG"],
        "increaseday_range": 21,
        "highday_range": 25,
        "lowday_range": 17,
        "sortType": 1,
        "tradeType": 1,
        "isReal": False,
        "startDate": None  # 默认不设置开始日期，使用全部数据
    }
    
    # 获取参数（优先使用POST请求中的参数）
    if request.method == 'POST':
        params = request.json or default_params
    else:
        params = default_params
    
    # 确保所有必要参数都存在
    for key, value in default_params.items():
        if key not in params:
            params[key] = value
    
    # 获取沪深300数据
    if params['ischeckhs300']:
        hs300_data = get_hs300_data(params['hs300_limit'], params['backCycle'], params.get('startDate'))
        params['hs300_data'] = hs300_data
    
    # 实例化策略并执行
    strategy = TradingStrategy(params)
    trade_log, cumulative_returns = strategy.execute()
    
    # 计算回撤和月度收益
    top_drawdowns = calculate_drawdowns(cumulative_returns)
    monthly_returns = calculate_monthly_returns(cumulative_returns)
    yearly_returns = calculate_yearly_returns(cumulative_returns)
    
    # 构建响应
    response = {
        'trade_log': trade_log,
        'cumulative_returns': cumulative_returns.to_dict('records'),
        'top_drawdowns': top_drawdowns,
        'monthly_returns': monthly_returns,
        'yearly_returns': yearly_returns
    }
    
    return jsonify(response)


